{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib widget\n",
    "import re\n",
    "import math\n",
    "import bisect\n",
    "import textwrap\n",
    "import numpy as np\n",
    "\n",
    "from pathlib import Path\n",
    "import mplcursors\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.patches import Rectangle, Patch\n",
    "\n",
    "kv_pattern = re.compile(r\";\\s*([\\w/.]+)\\s*=\\s*(['\\\"]?)([:,_\\w\\s/\\-\\\".']+)\\2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "KB = 1024\n",
    "# TPU Specs\n",
    "# local_mem_size: local memory size per npu\n",
    "# local_mem_banks: number of local memory banks\n",
    "# local_mem_size = local_mem_size // local_mem_banks\n",
    "bm1684x_spec = {\n",
    "    \"local_mem_size\": 256 * KB,\n",
    "    \"local_mem_banks\": 16 # 16 banks， local_mem_size = local_mem_size // local_mem_banks\n",
    "}\n",
    "\n",
    "bm1688_spec = {\n",
    "    \"local_mem_size\": 128 * KB,\n",
    "    \"local_mem_banks\": 16 # 16 banks， local_mem_size = local_mem_size // local_mem_banks\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 正则化\n",
    "def is_number_regex(s):\n",
    "    pattern = r\"^[-+]?(\\d+|\\d+\\.\\d*|\\.\\d+)([eE][-+]?\\d+)?$\"\n",
    "    return re.fullmatch(pattern, s) is not None\n",
    "\n",
    "def parse_dic(line, filter=None):\n",
    "    if filter is None:\n",
    "        filter = set()\n",
    "    elif isinstance(filter, str):\n",
    "        filter = set(filter.split(\",\"))\n",
    "    elif isinstance(filter, list):\n",
    "        filter = set(filter)\n",
    "\n",
    "    hit = True\n",
    "    for f in filter:\n",
    "        if f not in line:\n",
    "            hit = False\n",
    "            break\n",
    "    if not hit:\n",
    "        return None\n",
    "\n",
    "    ret = kv_pattern.findall(line)\n",
    "    dic = {}\n",
    "    for k, _, v in ret:\n",
    "        v = v.strip()\n",
    "        if is_number_regex(v):\n",
    "            dic[k] = int(v)\n",
    "        elif \"0x\" in v:\n",
    "            dic[k] = int(v, 16)\n",
    "        else:\n",
    "            dic[k] = v\n",
    "    return dic\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 解析log文件\n",
    "def dump_lmem_assign_result_with_setttings(log_file):\n",
    "    allocation_patches = []\n",
    "    allocation_settings = []\n",
    "    allocation = []\n",
    "    allocation_dict = {}\n",
    "    allocation_keys = [\"timestep_start\", \"timestep_end\", \"addr\", \"size\", \"op_type\", \"op_name\", \"lmem_type\", \"hold_in_lmem\", \"status\", \"allow_bank_conflict\", \"one_loop\", \"shape_secs\"]\n",
    "    allocation_setting = {}\n",
    "    allocation_setting_keys = [\"allow_bank_conflict\", \"shape_secs\"]\n",
    "    # allocation_setting_idx = 0\n",
    "    with open(log_file, 'r') as f:\n",
    "        for line in f:\n",
    "            dic = parse_dic(line, [\"; action = lmem_assign\", \"; tag = iteration_result\"])\n",
    "            if dic:\n",
    "                for key in allocation_keys:\n",
    "                    if key in dic:\n",
    "                        allocation_dict[key] = dic[key]\n",
    "                cur_setting = {}\n",
    "                for key in allocation_setting_keys:\n",
    "                    if key in dic:\n",
    "                        cur_setting[key] = dic[key]\n",
    "                if allocation_setting:\n",
    "                    if allocation_setting != cur_setting:\n",
    "                        allocation_settings.append(allocation_setting)\n",
    "                        allocation_patches.append(allocation)\n",
    "                        allocation_setting = {}\n",
    "                        allocation = []\n",
    "                    allocation.append(allocation_dict)\n",
    "                    allocation_dict = {}\n",
    "                else:\n",
    "                    allocation.append(allocation_dict)\n",
    "                    allocation_dict = {}\n",
    "                    allocation_setting = cur_setting\n",
    "        if allocation_setting:\n",
    "            allocation_settings.append(allocation_setting)\n",
    "            allocation_patches.append(allocation)\n",
    "    for settings in allocation_settings:\n",
    "        print(settings)\n",
    "    return allocation_settings, allocation_patches\n",
    "\n",
    "def update_allocation(allocation_patches):\n",
    "    allocation_patches_updated = []\n",
    "    for allocation in allocation_patches:\n",
    "        max_addr = 0\n",
    "        max_timestep = 0\n",
    "        allocation_success = []\n",
    "        allocation_failed = []\n",
    "        allocation_updated = []\n",
    "        for iter in allocation:\n",
    "            if iter[\"status\"] == \"success\":\n",
    "                allocation_success.append(iter)\n",
    "            elif iter[\"status\"] == \"failed\":\n",
    "                allocation_failed.append(iter)\n",
    "            max_timestep = max(max_timestep, iter[\"timestep_end\"])\n",
    "        # calculate max addr\n",
    "        for iter in allocation_success:\n",
    "            max_addr = max(max_addr, iter[\"addr\"] + iter[\"size\"])\n",
    "        # update failed allocation\n",
    "        for iter in allocation_failed:\n",
    "            if iter[\"status\"] == \"failed\":\n",
    "                iter[\"addr\"] = max_addr\n",
    "                max_addr += iter[\"size\"]\n",
    "        allocation_updated = allocation_success + allocation_failed\n",
    "        # add \"max_timestep\" to allocation\n",
    "        for iter in allocation_updated:\n",
    "            iter[\"max_timestep\"] = max_timestep\n",
    "        allocation_patches_updated.append(allocation_updated)\n",
    "    return allocation_patches_updated\n",
    "\n",
    "def find_allocation_idx(allocation_settings, settings):\n",
    "    allocation_idx = []\n",
    "    for idx, allocation in enumerate(allocation_settings):\n",
    "        if allocation == settings:\n",
    "            allocation_idx.append(idx)\n",
    "    print(f\"allocation idx: {allocation_idx}\")\n",
    "    return allocation_idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot memory allocation\n",
    "def plot_memory_allocation(allocations, allocation_setting, chip_spec=None, figsize=(12, 8), **kwargs):\n",
    "    # 自定义参数\n",
    "    timestep = kwargs.get(\"timestep\", None)\n",
    "    timestep_mode = True if timestep else False\n",
    "    timestep_mode_addr = 0\n",
    "    print(timestep)\n",
    "\n",
    "    get_lmem_per_timestep = kwargs.get(\"get_lmem_per_timestep\", False)\n",
    "\n",
    "    plt.rcParams.update({\n",
    "        'path.simplify': True,\n",
    "        'path.simplify_threshold': 1.0,\n",
    "        'agg.path.chunksize': 10000,\n",
    "    })\n",
    "\n",
    "    main_fig = plt.figure(figsize=figsize, dpi=100)\n",
    "    ax = plt.gca()\n",
    "    global current_annotation\n",
    "    current_annotation = None\n",
    "\n",
    "    all_times = [t for a in allocations for t in (a[\"timestep_start\"], a[\"timestep_end\"])]\n",
    "    total_timesteps = max(all_times) + 1 if all_times else 1\n",
    "\n",
    "    lmem_per_timestep = [0] * total_timesteps\n",
    "\n",
    "    colors = plt.cm.get_cmap('tab20', len(allocations))\n",
    "\n",
    "    max_addr = max([a[\"addr\"]+a[\"size\"] for a in allocations]) if allocations else 0\n",
    "    max_memory = bank_size = None\n",
    "    if chip_spec:\n",
    "        max_memory = chip_spec[\"local_mem_size\"]\n",
    "        bank_size = chip_spec[\"local_mem_size\"] // chip_spec[\"local_mem_banks\"]\n",
    "        y_max = max(max_memory, math.ceil(max_addr / bank_size) * bank_size)\n",
    "    else:\n",
    "        y_max = max_addr * 1.1\n",
    "\n",
    "    hatch_mapping = {\n",
    "        \"LMEM_ACTIVATION\": \"\",\n",
    "        \"LMEM_WEIGHT\": \"/////\",\n",
    "        \"LMEM_OPERATION\": \"***\",\n",
    "    }\n",
    "\n",
    "    rects = []\n",
    "    legend_elements = []\n",
    "\n",
    "    for idx, allocation in enumerate(allocations):\n",
    "        s = allocation[\"timestep_start\"]\n",
    "        e = allocation[\"timestep_end\"]\n",
    "        addr = allocation[\"addr\"]\n",
    "        size = allocation[\"size\"]\n",
    "        op_type = allocation[\"op_type\"]\n",
    "        op_name = allocation[\"op_name\"]\n",
    "        lmem_type = allocation[\"lmem_type\"]\n",
    "        hold_in_lmem = allocation[\"hold_in_lmem\"]\n",
    "        one_loop = allocation[\"one_loop\"]\n",
    "        status = allocation[\"status\"]\n",
    "        color = colors(idx)\n",
    "        edgecolor = 'black' if status == \"success\" else 'red'\n",
    "\n",
    "        if not one_loop and hold_in_lmem:\n",
    "            for t in range(total_timesteps):\n",
    "                lmem_per_timestep[t] += size\n",
    "        else:\n",
    "            if s <= e:\n",
    "                for t in range(s, e+1):\n",
    "                    lmem_per_timestep[t] += size\n",
    "            else:\n",
    "                for t in range(s, total_timesteps):\n",
    "                    lmem_per_timestep[t] += size\n",
    "                for t in range(0, e+1):\n",
    "                    lmem_per_timestep[t] += size\n",
    "\n",
    "        time_ranges = []\n",
    "        if s <= e:\n",
    "            time_ranges.append((s, e-s+1))\n",
    "        else:\n",
    "            time_ranges.append((s, total_timesteps-s))\n",
    "            time_ranges.append((0, e+1))\n",
    "        if hold_in_lmem and not one_loop and not timestep:\n",
    "            time_ranges.append((0, total_timesteps))\n",
    "        if timestep:\n",
    "            res = []\n",
    "            for t_start, t_duration in time_ranges:\n",
    "                if t_start <= timestep < t_start + t_duration:\n",
    "                    res = [(timestep, 1)]\n",
    "                    break\n",
    "            time_ranges = res\n",
    "            if res:\n",
    "                addr = timestep_mode_addr\n",
    "                timestep_mode_addr += size\n",
    "\n",
    "        for t_start, t_duration in time_ranges:\n",
    "            rect = Rectangle(\n",
    "                (t_start, addr),\n",
    "                t_duration,\n",
    "                size,\n",
    "                facecolor=color,\n",
    "                edgecolor=edgecolor,\n",
    "                linewidth=0.5,\n",
    "                alpha=0.7,\n",
    "                picker=True,\n",
    "                snap=True,\n",
    "                hatch=hatch_mapping.get(lmem_type, \"\"),\n",
    "            )\n",
    "\n",
    "            rect.block_info = {\n",
    "                \"block_id\": idx,\n",
    "                \"global_start\": s,\n",
    "                \"global_end\": e,\n",
    "                \"segment_start\": t_start,\n",
    "                \"segment_duration\": t_duration,\n",
    "                \"address\": addr,\n",
    "                \"size\": size,\n",
    "                \"color\": color,\n",
    "                \"op_type\": op_type,\n",
    "                \"op_name\": op_name,\n",
    "                \"lmem_type\": lmem_type,\n",
    "                \"hold_in_lmem\": \"True\" if hold_in_lmem else \"False\"\n",
    "            }\n",
    "            ax.add_patch(rect)\n",
    "            rects.append(rect)\n",
    "\n",
    "        legend_elements.append(Patch(\n",
    "            facecolor=color,\n",
    "            label=f'Block {idx}: {size}B @ {addr}',\n",
    "        ))\n",
    "\n",
    "    cursor = mplcursors.cursor(\n",
    "        rects,\n",
    "        hover=False,\n",
    "        highlight=False\n",
    "    )\n",
    "\n",
    "    @cursor.connect(\"add\")\n",
    "    def on_hover(sel):\n",
    "        global current_annotation\n",
    "        sel.annotation.set_visible(False)\n",
    "\n",
    "        if current_annotation:\n",
    "            try:\n",
    "                current_annotation.remove()\n",
    "            except (ValueError, AttributeError):\n",
    "                pass\n",
    "            current_annotation = None\n",
    "\n",
    "        rect = sel.artist\n",
    "        info = rect.block_info\n",
    "\n",
    "        rgb = matplotlib.colors.to_rgb(info['color'])\n",
    "        luminance = 0.299*rgb[0] + 0.587*rgb[1] + 0.114*rgb[2]\n",
    "        text_color = 'black' if luminance > 0.5 else 'white'\n",
    "        xtext = 15 if info['segment_start'] < total_timesteps // 2 else -200\n",
    "        ytext = 15 if info['address'] < y_max // 2 else -100\n",
    "\n",
    "        current_annotation = ax.annotate(\n",
    "            text=build_tooltip(info),\n",
    "            xy=(info['segment_start'], info['address']),\n",
    "            xytext=(xtext, ytext),\n",
    "            textcoords='offset points',\n",
    "            bbox=dict(\n",
    "                boxstyle=\"round,pad=0.3\",\n",
    "                facecolor=matplotlib.colors.to_rgba(info['color'], 0.9),\n",
    "                edgecolor='black',\n",
    "                linewidth=0.5\n",
    "            ),\n",
    "            arrowprops=dict(\n",
    "                arrowstyle=\"->\",\n",
    "                connectionstyle=\"arc3,rad=0.3\",\n",
    "                color='black'\n",
    "            ),\n",
    "            fontsize=9,\n",
    "            color=text_color\n",
    "        )\n",
    "        main_fig.canvas.draw_idle()\n",
    "\n",
    "    def build_tooltip(info):\n",
    "        return (\n",
    "            f\"Block ID: {info['block_id']}\\n\"\n",
    "            f\"Timestep: {info['global_start']}→{info['global_end']}\\n\"\n",
    "            f\"Address: {info['address']}-{info['address']+info['size']}\\n\"\n",
    "            f\"Size: {info['size']}B ({info['size']/1024:.1f} KB)\\n\"\n",
    "            f\"Op Type: {info['op_type']}\\n\"\n",
    "            f\"Op Name: {textwrap.fill(info['op_name'], width=50)}\\n\"\n",
    "            f\"LMEM Type: {info['lmem_type']}\\n\"\n",
    "            f\"Hold in LMEM: {info['hold_in_lmem']}\"\n",
    "        )\n",
    "\n",
    "    def on_leave(event):\n",
    "        if current_annotation and current_annotation.get_visible():\n",
    "            current_annotation.remove()\n",
    "            main_fig.canvas.draw_idle()\n",
    "\n",
    "    main_fig.canvas.mpl_connect('axes_leave_event', on_leave)\n",
    "\n",
    "    for ts in range(total_timesteps + 1):\n",
    "        ax.axvline(ts, color='gray', linestyle='--', alpha=0.3)\n",
    "\n",
    "    if max_memory is not None:\n",
    "        ax.axhline(y=max_memory,\n",
    "                  color='black',\n",
    "                  linestyle='--',\n",
    "                  linewidth=2)\n",
    "\n",
    "    ax.set_ylim(0, y_max)\n",
    "    ax.set_ylabel(\"Memory Address\")\n",
    "\n",
    "    if bank_size:\n",
    "        num_banks = y_max // bank_size\n",
    "        ax.set_yticks([i*bank_size for i in range(num_banks+1)])\n",
    "\n",
    "        for bank in range(1, num_banks):\n",
    "            ax.axhline(y=bank*bank_size,\n",
    "                      color='gray',\n",
    "                      linestyle='--',\n",
    "                      linewidth=1.5,\n",
    "                      alpha=0.5)\n",
    "\n",
    "        ax2 = ax.twinx()\n",
    "        ax2.set_ylim(ax.get_ylim())\n",
    "        ax2.set_yticks([(i+0.5)*bank_size for i in range(num_banks)])\n",
    "        ax2.set_yticklabels([f\"Bank {i}\" for i in range(num_banks)],\n",
    "                            fontdict={\n",
    "                                'fontsize': 8,\n",
    "                                'color': 'gray',\n",
    "                                'alpha': 0.7\n",
    "                            })\n",
    "        ax2.set_ylabel('Memory Banks')\n",
    "\n",
    "    spec_text = f\"SPEC: LmemSize(Byte)={max_memory or 'N/A'}, BankSize(Byte)={bank_size or 'N/A'}\"\n",
    "\n",
    "    setting_text = \"\\nSETTINGS: \"\n",
    "    for k, v in allocation_setting.items():\n",
    "        setting_text += f\"{k}={v}, \"\n",
    "    setting_text = setting_text[:-2]\n",
    "    ax.set_xlabel(\"Timestep\")\n",
    "    ax.set_title(f\"TPU Memory Allocation \\n{spec_text}{setting_text}\", pad=30)\n",
    "    ax.set_xticks([i + 0.5 for i in range(total_timesteps)])\n",
    "    ax.set_xticklabels([str(i) for i in range(total_timesteps)])\n",
    "    ax.set_xlim(0, total_timesteps)\n",
    "\n",
    "    if max_memory:\n",
    "        legend_elements.append(Patch(\n",
    "            facecolor='white',\n",
    "            edgecolor='black',\n",
    "            linestyle='--',\n",
    "            linewidth=2,\n",
    "            label=f'Max Memory ({max_memory}B)'\n",
    "        ))\n",
    "    for lmem_type, hatch in hatch_mapping.items():\n",
    "        legend_elements.append(Patch(\n",
    "            facecolor='white',\n",
    "            edgecolor='black',\n",
    "            linestyle='-',\n",
    "            hatch=hatch,\n",
    "            linewidth=0.5,\n",
    "            label=f'{lmem_type}'\n",
    "        ))\n",
    "\n",
    "    ax.legend(handles=legend_elements,\n",
    "             bbox_to_anchor=(-0.1, 1.2, 0, 0),\n",
    "             loc='upper right',\n",
    "             borderaxespad=3,\n",
    "             fontsize=7,\n",
    "    )\n",
    "\n",
    "    def on_resize(event):\n",
    "        ax = event.canvas.figure.axes[0]\n",
    "        fig_width = event.width / 100\n",
    "        legend = ax.legend(prop={'size': fig_width * 0.8})\n",
    "        event.canvas.draw()\n",
    "\n",
    "    main_fig.canvas.mpl_connect('resize_event', on_resize)\n",
    "\n",
    "    plt.tight_layout()\n",
    "\n",
    "    hist_fig = None\n",
    "    if get_lmem_per_timestep:\n",
    "        hist_fig = plt.figure(figsize=(figsize[0], figsize[1]//2))\n",
    "        hist_ax= hist_fig.gca()\n",
    "\n",
    "        timesteps = list(range(len(lmem_per_timestep)))\n",
    "        bars = hist_ax.bar(timesteps, lmem_per_timestep, width=0.8,\n",
    "                          edgecolor='black', alpha=0.7)\n",
    "\n",
    "        for bar in bars:\n",
    "            height = bar.get_height()\n",
    "            if height > 0:\n",
    "                hist_ax.text(bar.get_x() + bar.get_width()/2., height,\n",
    "                            f'{height}B',\n",
    "                            ha='center', va='bottom', fontsize=8)\n",
    "\n",
    "        hist_ax.set_xlabel(\"Timestep\")\n",
    "        hist_ax.set_ylabel(\"Memory Usage (Bytes)\")\n",
    "        hist_ax.set_title(\"Memory Usage Per Timestep\", pad=15)\n",
    "        hist_ax.set_xticks(timesteps)\n",
    "        hist_ax.grid(axis='y', linestyle='--', alpha=0.7)\n",
    "        hist_fig.tight_layout()\n",
    "\n",
    "    return (main_fig, hist_fig) if get_lmem_per_timestep else main_fig\n",
    "\n",
    "def plot_memory_allocation_with_idx(allocation_patches, allocation_settings, allocation_idx, chip_spec=None, figsize=(12, 8), **kwargs):\n",
    "    timestep = kwargs.get(\"timestep\", None)\n",
    "    get_lmem_per_timestep = kwargs.get(\"get_lmem_per_timestep\", False)\n",
    "    for idx in allocation_idx:\n",
    "        allocations = allocation_patches[idx]\n",
    "        setting = allocation_settings[idx]\n",
    "        print(f\"allocation idx: {idx}\\nsettings: {setting}\")\n",
    "        if get_lmem_per_timestep:\n",
    "            fig, hist_fig = plot_memory_allocation(allocations, setting, chip_spec, figsize, timestep=timestep, get_lmem_per_timestep=True)\n",
    "        else:\n",
    "            fig = plot_memory_allocation(allocations, setting, chip_spec, figsize, timestep=timestep)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "log_file = \"lmem_assign_result.log\"\n",
    "allocation_settings, allocation_patches = dump_lmem_assign_result_with_setttings(log_file)\n",
    "allocation_patches_updated = update_allocation(allocation_patches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set according to info above\n",
    "setting = {\n",
    "    \"allow_bank_conflict\": False, # manually setting\n",
    "    \"shape_secs\": \"8,7,1,1,1\"     # manually setting\n",
    "}\n",
    "allocation_idx = find_allocation_idx(allocation_settings, setting)\n",
    "plot_memory_allocation_with_idx(allocation_patches_updated, allocation_settings, allocation_idx, bm1688_spec, (12, 8), get_lmem_per_timestep=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_memory_allocation_with_idx(allocation_patches_updated, allocation_settings, allocation_idx, bm1688_spec, (12, 8), timestep=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
